/**
  ETFAna project, Anyang Normal University && IMP-CAS
  \class ETFPostMagTracker
  \brief trackers for tracking downstream the dipole magnet
  \author SUN Yazhou, asia.rabbit@163.com
  \since 2021-11-02
  \date 2024-01-09 last modified
  \attention
  changelog
  <table>
  <tr>  <th>Date         <th>Author      <th>Description                   </tr>
  <tr>  <td>2021-11-02   <td>Asia Sun    <td>file created                  </tr>
  </table>

  \copyright Copyright (c) 2021-2024 Anyang Normal U. && IMP-CAS with LGPLv3 LICENSE
*/

#include <iostream>
#include <TClonesArray.h>
#include "ETFRun.h"
#include "ETFParaManager.h"
#include "ETFPostMagTracker.h"
#include "ETFDetectorManager.h"
#include "ETFVDC.h"
#include "ETFMath.h"
#include "ETFTrack.h"
#include "ETF3DTrack.h"
#include "ETFTOFWall.h"
#include "ETFDCHit.h"
#include "ETFPlaHit.h"
#include "ETFMsg.h"
#include "ETFYaml.h"
#include "ETFTimer.h"

using std::cout;
using std::endl;
using std::flush;

#define er ETFMsg::Error

ETFPostMagTracker::ETFPostMagTracker(const char *name) : ETFTracker(name){}

ETFPostMagTracker::ETFPostMagTracker(const char *name, const vector<string> &dcs,
const string &tofw, const string &t0) : ETFTracker(name, dcs, t0), fTOFWall(0),
fPhiAvrg(-999.){
  if(fBroken) return;
  if("" != tofw){ // use tofw as the timing det
    fTOFWall = dynamic_cast<ETFTOFWall *>(fdm->GetDetector(tofw));
    // nstripStray: limit of distance that a trk lies astray from the center of
    // its corresponding fired TOFWall strip
    const ETFYaml &nss = (*fpm->GetMisPara())[name]["nstripStray"];
    if(nss) fTOFWall->SetNStripStrayBound(nss[0].as<double>(), nss[1].as<double>());
  } // end if
  if(!fTOFWall && !fTimer[0]){
    // er("ETFPostMagTracker", "ctor: no timing det is specified");
    fBroken = true; return;
  } // end if
  if(fTOFWall && fTimer[0]){
    er("ETFPostMagTracker", "ctor: tofw and t0 both specified");
    fBroken = true; return;
  } // end if

  const char xuv[] = {'X', 'U', 'V'};
  char tname[128];
  for(int i = 3; i--;){
    fTracks[i] = new TClonesArray("ETFTrack");
    sprintf(tname, "ETFTrack%c", xuv[i]);
    fTracks[i]->SetName(tname);
    fTracksOut[i] = new TClonesArray(*fTracks[i]);
    sprintf(tname, "%s%c", name, xuv[i]);
    fTracksOut[i]->SetName(tname);
  } // end if
  f3DTracks = new TClonesArray("ETF3DTrack");
  sprintf(tname, "%s3D", name);
  f3DTracks->SetName(tname);

  for(auto d : fDCs) fPhiAvrg += d->GetPhi();
  fPhiAvrg /= fDCs.size();
} // end ctor

ETFPostMagTracker::~ETFPostMagTracker(){}

bool ETFPostMagTracker::SetDriftDistance(ETFVTrack *trk){
  return fDCs[0]->SetDriftDistance(trk);
} // end member function SetDriftDistance

// update the hit position of the tracks in the tracker
void ETFPostMagTracker::UpdateHitPos(){
  ETFTracker::UpdateHitPos();
  if(fTOFWall) UpdateHitPos(fTOFWall);
  if(fTimer[0]) UpdateHitPos(fTimer[0]);
} // end member function UpdateHitPos

// X tracking
void ETFPostMagTracker::Map(){
  if(fBroken) return;
  if(fTOFWall && fTOFWall->GetNFiredStrip() <= 0) return;
  if(fTimer[0] && fTimer[0]->GetFiredStatus() != 4) return;
  for(auto &dc : fDCs) dc->Filter();
  
  Map(TrackType::kX);

  UpdateHitPos();
  if(fTOFWall){
    int i = 0;
    for(auto b : *fTracksOut[0]){
      ETFTrack *t = static_cast<ETFTrack *>(b);
      fTOFWall->GetStrip(t->GetFiredStripId())->SetTrackId(i++);
    } // end for
  } // end if
} // end member functuion Map

// 3D tracking
void ETFPostMagTracker::Map3D(){
  if(fBroken) return;
  if(fTOFWall && fTOFWall->GetNFiredStrip() <= 0) return;
  if(fTimer[0] && fTimer[0]->GetFiredStatus() != 4) return;
  
  Map(TrackType::kX);
  if(!GetNofTracks(TrackType::kX)) return;

  Map(TrackType::kU);
  if(!GetNofTracks(TrackType::kU)) return;

  Map(TrackType::kV);
  if(!GetNofTracks(TrackType::kV)) return;

  TrackMerge();
  UpdateHitPos();
} // end member functuion Map3D

// called by Map, to test if a hit combination is eligible
bool ETFPostMagTracker::TestCombination(TrackType type){
  const int n = ETFTrack::EnoughHit(fHits);
  if(!n) return false;

  vector<double> z, x, dx;
  z.resize(n); x.resize(n); dx.resize(n);
  dx.assign(n, -1.); // dx not unkown, so set to minus
  int i = 0;
  for(auto &p : fHits) if(p.second){ p.second->GetProjZX(z[i], x[i]); i++; }
  double k = -999., b = -999., d2 = 1.e200;
  ETFMath::LinearFit(z.data(), x.data(), dx.data(), n, b, k, d2);
  if(d2 > D2THREPD*n) return false;

  const short t = short(type);
  ETFTrack *trk = static_cast<ETFTrack *>(ConstructedAt(fTracks[t]));
  if(!trk) return false; // probably the clonesarray is full (>= 1000)
  // assign the drift time start and fired strip id
  trk->Initialize(); trk->SetDsquare(d2);
  trk->SetK(k); trk->SetB(b);
  trk->SetId(fTracks[t]->GetEntriesFast()-1, t); // set name&title of the track
  for(auto &p : fHits) if(p.second) trk->AddHit(p.second);
  if(TrackType::kX != type) return true; // the following tests are only for X tracks

  // valid TOFWall strip not found
  if(fTOFWall){
    if(!fTOFWall->AssignDriftTimeStart(trk)){
      fTracks[t]->Remove(trk);
      return false;
    } // end inner if
  } // end outer if
  else if(fTimer[0]){
    trk->SetDriftTimeStart(fTimer[0]->GetTime());
    trk->SetFiredStripId(0); // just assign 0 here
    trk->SetNStripStray(0.); // and here
  } // end else if
  // here whether drift time is within range is checked
  if(!SetDriftDistance(trk)){ fTracks[t]->Remove(trk); return false; }
  // fit the track with drift distance, and assign the fitting residues
  // using space-time relation specific to post-mag dc, i.e. ETFDC
  // if the fitting result fails the user requirements, returns false
  if(fChiPD > 0.) trk->SetChiPD(fChiPD); // user-input chiPD
  if(!trk->Fit()){ fTracks[t]->Remove(trk); return false; }
  // compare the current track with the existing tracks for compatibility check
  // remove the track if incompatibility is encountered
  return Compare(trk, type);
} // end member function TestCombination

// compare trk with all the stored tracks of the same type
// if trk is considered to belong to a different particle
// it would then be accepted and added to the list of tracks
// return true if the new track is accepted
// \param trk1: newtrk; trk0: oldtrk
bool ETFPostMagTracker::Compare(ETFVTrack *trk1, ETFVTrack *trk0, TrackType type){
  if(TrackType::kX != type) er("ETFPostMagTracker", "Compare: intended only for X tracks");
  TClonesArray *c = fTracks[short(type)];
  const int cmp = Compare(trk1, trk0);
  if(2 == cmp) c->Remove(trk0);
  if(0 == cmp) c->Remove(trk1);
  return cmp;
} // end member function Compare

// the same as the TrackType version, but using track array c instead of type
// \retval 0: newtrk is bad; 1: two tracks are different (keep both); 2: oldtrk is bad
int ETFPostMagTracker::Compare(ETFVTrack *newtrk, ETFVTrack *oldtrk){
  ETFTrack *trk0 = static_cast<ETFTrack *>(oldtrk);
  ETFTrack *trk1 = static_cast<ETFTrack *>(newtrk);
  int nstripDeviation = abs(trk0->GetFiredStripId() - trk1->GetFiredStripId());
  // TODO: user input configuration expected here
  static const int NL = fDCs.size() * 2;
  static const int VICINITY = 3; // min nof cells 2 seperate tracks can be apart from for
  static const int STRIPTOL = 1; // maximal TOFWall strip deviation tolerance
  // number of points (hits) in the new and old tracks
  int np0 = trk0->GetNofHits(), np1 = trk1->GetNofHits();
  int vicinityViaCnt = 0; // count of vicinity vialations
  int nCom = 0; // nof layers where trk0 and trk1 both have hits
  for(int i = 0; i < NL; i++) if((*trk0)[i] && (*trk1)[i]) nCom++;
  // -------------- the new track has less hits than the old track -------------//
  if(np1 < np0){
    for(int i = 0; i < NL; i++)
      if((*trk0)[i] && (*trk1)[i] && abs((*trk0)[i]->GetNu()-(*trk1)[i]->GetNu()) > VICINITY)
        // if(++vicinityViaCnt == np1) return 1; // the two tracks are different
        if(++vicinityViaCnt == nCom) return 1; // the two tracks are different
    // TOFWall strips are checked if >= 1 pair of hits are too close between the 2 tracks
    if(nstripDeviation <= STRIPTOL) return 0; // new track is part of the old track
    return 1; // TOFW strips are far apart enough -- we can as well accept the new track
  } // end if(np1 < np0)
  // -------------- the new track has more hits than the old track -------------//
  if(np1 > np0){
    for(int i = 0; i < NL; i++)
      if((*trk0)[i] && (*trk1)[i] && abs((*trk0)[i]->GetNu()-(*trk1)[i]->GetNu()) > VICINITY)
        // if(++vicinityViaCnt == np0) return 1; // the two tracks are different
        if(++vicinityViaCnt == nCom) return 1; // the two tracks are different
    // TOFWall strips are checked if >= 1 pair of hits are too close between the 2 tracks
    if(nstripDeviation <= STRIPTOL) return 2;
    return 1; // TOFW strips are far apart enough -- we can as well accept the new track
  } // end if(np1 > np0)
  // -------------- the new track has equal hits to the old track -------------//
  if(np1 == np0){
    for(int i = 0; i < NL; i++)
      if((*trk0)[i] && (*trk1)[i] && abs((*trk0)[i]->GetNu()-(*trk1)[i]->GetNu()) > VICINITY)
        // if(++vicinityViaCnt == np0) return 1; // the two tracks are different
        if(++vicinityViaCnt == nCom) return 1; // the two tracks are different
    // TOFWall strips are checked if >= 1 pair of hits are too close between the 2 tracks
    if(nstripDeviation <= STRIPTOL){ // the fitting chis are compared for further selection
      // newtrk has better fitting quality, so acceptable
      if(trk1->GetChi2() < trk0->GetChi2()) return 2;
      return 0; // old track beats the new track in tracking quality
    } // end if
    return 1; // TOFW strips are far apart enough -- we can as well accept the new track
  } // end if(np1 > np0)

  return 1; // never will reach here, just for suppressing compiler warnings
} // end member function Compare(c)

// match XUV track projections to 3D tracks
void ETFPostMagTracker::TrackMerge(){
  static constexpr TrackType t[3] = {TrackType::kX, TrackType::kU, TrackType::kV};
  int n[3]{};
  for(int i = 3; i--;) if(!(n[i] = fTracksOut[i]->GetEntriesFast())) return;
  // static int ntrk[3]{}; // DEBUG
  // for(int i = 3; i--;) ntrk[i] += fTracksOut[i]->GetEntriesFast(); // DEBUG
  // cout << "ntrkX-U-V:" << ntrk[0] << "-" << ntrk[1] << "-" << ntrk[2] << "-" << endl << endl; // DEBUG


  // loop over every possible combination of (X,U,V) tracks to match them //
  int id = -1; // 3D track identifier
  double k2, b2;
  for(int i = 0; i < n[0]; i++){ // loop over track X
    ETFTrack *x = static_cast<ETFTrack *>(GetTrack(i, t[0]));
    bool isMatched = false; // whether the current X track has found its companies
    for(int j = 0; j < n[1]; j++){ // loop over track U
      ETFTrack *u = static_cast<ETFTrack *>(GetTrack(j, t[1]));
      int id0 = id + 1; if(isMatched) id0 = id; // the current 3D track id
      // one projection cannot be matched with multiple 3D tracks
      if(u->Get3DId() != -1 && u->Get3DId() < id0) continue; // owned by previous X tracks
      for(int k = 0; k < n[2]; k++){ // loop over track V
        ETFTrack *v = static_cast<ETFTrack *>(GetTrack(k, t[2]));
        if(v->Get3DId() != -1 && v->Get3DId() < id0) continue; // owned by previous X tracks
        double xMiss3D[3]{};
        if(Compatible(x, u, v, k2, b2, xMiss3D)){ // get a match
          const int sid = x->GetFiredStripId();
          const double nstray = x->GetNStripStray();
          const double tof = fTOFWall->GetStrip(sid)->GetTime();
          // assign fired strip of X to U and V, and do the fitting //
          u->SetFiredStripId(sid); u->SetNStripStray(nstray);
          u->SetDriftTimeStart(tof);
          if(!SetDriftDistance(u) || !u->Fit()) break; // break to next u
          // here's V's turn
          v->SetFiredStripId(sid); v->SetNStripStray(nstray);
          v->SetDriftTimeStart(tof);
          if(!SetDriftDistance(v) || !v->Fit()) continue; // continue to next v
          // acknowledge the newly found 3D track
          if(!isMatched){ id++; isMatched = true; } // isMatched only set once for each X
          x->Set3DId(id); u->Set3DId(id); v->Set3DId(id);
        } // end if
      } // end for over track V
    } // end for over track U
  } // end for over track X
  if(-1 == id) return; //< no 3d track has been found

  // summarize the results //
  // clean eliminated tracks and filter incompatible ones
  id += 1;
  CleanUp(TrackType::kU, id); // clean U tracks
  CleanUp(TrackType::kV, id); // clean V tracks
  // add matched tracks to f3DTracks
  for(int i = 0; i < id; i++){ // loop over 3D tracks
    ETF3DTrack *trk3d = static_cast<ETF3DTrack *>(f3DTracks->ConstructedAt(i));
    trk3d->Initialize();
    for(int j = 0; j < 3; j++){ // loop over XUV
      for(auto p : *fTracksOut[j]){ // loop over all the tracks of tracktype t[j]
        ETFTrack *trk = static_cast<ETFTrack *>(p);
        if(i == trk->Get3DId()){ trk3d->SetTrack(t[j], trk); continue; }
      } // end for over tracks
    } // end for over XUV
    // complete other 3D track info //
    const auto &m = trk3d->GetTrackMap();
    ETFVTrack *x = m.at(t[0]), *u = m.at(t[1]), *v = m.at(t[2]);
    double xMiss3D[3]{};
    Compatible(x, u, v, k2, b2, xMiss3D);
    trk3d->SetK1B1(x->GetK(), x->GetB());
    trk3d->SetK2B2(k2, b2);
    trk3d->SetXMiss3D(xMiss3D);
    trk3d->SetId(i);
    trk3d->AssignTOT();
    u->SetHitFiredStatus(2);
    v->SetHitFiredStatus(2);
  } // end for over 3D tracks
} // end member function TrackMerge

// clean eliminated tracks and filter incompatible ones
void ETFPostMagTracker::CleanUp(TrackType type, int n){
  const short t = short(type);
  TClonesArray *c = fTracks[t];
  for(int i = 0; i < n; i++){
    ETFTrack *tm = nullptr; // the optimal track
    const int ntrk = GetNofTracks(type);
    // `compare()` all the tracks with the same 3DId (i) to select the best one //
    for(int j = 0; j < ntrk; j++){
      ETFTrack *tt = static_cast<ETFTrack *>(GetTrack(j, type));
      if(tt->Get3DId() != i) continue; // tracks with 3DId == -1 would be eliminated later
      if(!tm) tm = tt;
      if(tt != tm){
        int cmp = Compare(tt, tm);
        if(1 == cmp){ // shit happens
          if(tm->GetChi() > tt->GetChi()) cmp = 2; // tm is bad
          else cmp = 0; // tt is bad
        } // end if
        if(0 == cmp) tt->Set3DId(-1); // tt is defeated
        if(2 == cmp){ // tm is defeated
          tm->Set3DId(-1);
          tm = tt;
        } // end if(2)
      } // end if(tt != tm)
    } // end for over j
  } // end for over i

  // remove the unmatched and defeated tracks
  for(auto p : *c) if(-1 == static_cast<ETFTrack *>(p)->Get3DId()) c->Remove(p);
} // end member function CleanUp

void ETFPostMagTracker::Clean(TClonesArray *cIn, TClonesArray *cOut){
  cOut->Clear();
  int i = 0;
  for(const auto &p : *cIn){ // this scoped for skips nullptrs automatically
    ETFTrack *t = static_cast<ETFTrack *>(cOut->ConstructedAt(i++));
    t->copy(static_cast<ETFVTrack *>(p));
    // only assign for X trks -- U and V are done in the track merge
    if(t->GetFiredStripId() >= 0){
      t->SetHitFiredStatus(2);
      // update drift time in case that the hits are altered by
      // other (defeated) tracks with different TOFW strips
      if(t->GetFitStatus() != FitStatus::kUNFITTED){
        t->UpdateDriftTime();
        if(!SetDriftDistance(t) || !t->Fit()) er(GetName(), "Clean: Shit happens");
      } // end if
      t->UpdateChi(); // this is the final chi
    } // end if
  } // end for
} // end member function Clean

// to tell if (X,U,V) combination is from the same 3D track
// if so, return l: y=k2*z+b2 and the inconsistency xMiss3D[3] in the 3 DC central planes
bool ETFPostMagTracker::Compatible(ETFVTrack *x, ETFVTrack *u, ETFVTrack *v,
  double &k2, double &b2, double *xMiss3D){
  return Compatible(x->GetK(), x->GetB(), u->GetK(), u->GetB(), v->GetK(), v->GetB(),
    k2, b2, xMiss3D);
} // end member function Compatible

// the same as the one taking ETFVTrack parameters
bool ETFPostMagTracker::Compatible(double k, double b, double ku, double bu,
  double kv, double bv, double &k2, double &b2, double *xMiss3D){
    const double &phi = fPhiAvrg;
  	double k1t = ETFMath::kUV_X(phi, ku, kv);
  	double b1t = ETFMath::bUV_X(phi, ku, kv, bu, bv);
  	double z, x, xt;

  	for(int i = 0; i < 3; i++){
  		z = fDCs[i]->GetZ(); // the central Z of DC-i
  		xt = k1t*z+b1t; x = k*z+b;
  		xMiss3D[i] = x-xt;
  		// test coincidence in every MWDC
      // TODO: user input configuration expected here
      static const double WIN3D = 10.; // 10, in mm
      // cout << "ku: " << ku << " kv: " << kv << endl; // DEBUG
      // cout << "k: " << k << " k1t: " << k1t << endl; // DEBUG
      // getchar(); // DEBUG
  		if(fabs(x-xt) > WIN3D) return false;
  	} // end for over i

  	k2 = ETFMath::kUV_Y(phi, ku, kv);
  	b2 = ETFMath::bUV_Y(phi, ku, kv, bu, bv);
  	return true;
} // end member function Compatible

ETFTOFWall *ETFPostMagTracker::GetTOFWall() const{
  if(!fTOFWall) er("ETFPostMagTracker", "GetTOFWall: fTOWFall is nullptr");
  return fTOFWall;
} // end member function GetTOFWall
